"""
智能体基础接口
定义所有白方智能体必须实现的核心接口

作者：wjttdbx
版本：3.0 - 重构版本
"""
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Tuple, Optional
import numpy as np

# 避免循环导入，使用 TYPE_CHECKING
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from sh15.core.units import WhiteBoat, WhiteDrone, BlackBoat


class AgentAction:
    """
    智能体动作数据结构
    
    封装智能体生成的动作信息，提供类型安全和易于扩展的接口
    
    - 无人艇动作: [速度, 方向, lock_id]
    - 无人机动作: [速度, 方向, drone_action]
    """
    
    def __init__(self):
        self.boat_actions = np.zeros((5, 3))  
        # 5艘船，每艘3个动作: [速度, 方向, lock_id]
        self.drone_actions = np.zeros((0, 3))  
        # 动态数量的无人机，每架3个动作: [速度, 方向, drone_action]
        
    
    def set_boat_action(self, boat_id: int, speed: float, heading: float, lock_id: int):
        """设置无人艇动作: [速度, 方向, lock_id]"""
        if 0 <= boat_id < 5:
            self.boat_actions[boat_id] = [speed, heading, lock_id]
    
    def set_drone_actions(self, drone_actions: np.ndarray):
        """设置无人机动作: [速度, 方向, drone_action]"""
        self.drone_actions = drone_actions.copy()
    
    def get_boat_actions(self) -> np.ndarray:
        """获取无人艇动作数组: [速度, 方向, lock_id]"""
        return self.boat_actions
    
    def get_drone_actions(self) -> np.ndarray:
        """获取无人机动作数组: [速度, 方向, drone_action]"""
        return self.drone_actions
    


class AgentObservation:
    """
    智能体观测数据结构
    
    标准化智能体接收的观测信息
    """
    
    def __init__(self, 
                 white_boats: List['WhiteBoat'], 
                 all_drones: List['WhiteDrone'],
                 black_boats: List['BlackBoat'], 
                 current_time: float):
        self.white_boats = white_boats
        self.all_drones = all_drones
        self.black_boats = black_boats  # 仅包含探测到的黑方目标
        self.current_time = current_time
        
        # 缓存的计算结果
        self._cached_detected_targets = None
        self._cached_threat_analysis = None
        self._cached_single_detected_targets = None

    def get_single_detected_targets(self, drone: 'WhiteDrone') -> List['BlackBoat']:
        """获取单个无人机探测到的黑方目标"""
        if self._cached_single_detected_targets is None:
            detected_ids = set()
            
            # 收集无人机探测到的目标
            for target in drone.detected_targets:
                if target.active:
                    detected_ids.add(target.id)
            
            self._cached_single_detected_targets = [
                boat for boat in self.black_boats 
                if boat.id in detected_ids and boat.active
            ]
        
        return self._cached_single_detected_targets


    def get_detected_targets(self) -> List['BlackBoat']:
        """获取所有探测到的黑方目标（包括无人船和无人机探测到的目标）"""
        if self._cached_detected_targets is None:
            detected_ids = set()
            
            # 收集白方无人船探测到的目标
            for boat in self.white_boats:
                if boat.active:
                    for target in boat.detected_targets:
                        if target.active:
                            detected_ids.add(target.id)
            
            # 收集白方无人机探测到的目标
            for drone in self.all_drones:
                if drone.active:
                    for target in drone.detected_targets:
                        if target.active:
                            detected_ids.add(target.id)
            
            self._cached_detected_targets = [
                boat for boat in self.black_boats 
                if boat.id in detected_ids and boat.active
            ]
        
        return self._cached_detected_targets
    
    def get_active_boats(self) -> List['WhiteBoat']:
        """获取活跃的白方无人艇"""
        return [boat for boat in self.white_boats if boat.active]
    
    def get_active_drones(self) -> List['WhiteDrone']:
        """获取活跃的无人机"""
        return [drone for drone in self.all_drones if drone.active]
    
    def get_ready_drones(self) -> int:
        """获取准备就绪的无人机数量"""
        ready_count = 0
        for drone in self.all_drones:
            if drone.state == "ready" and drone.docked_on_boat_id is not None:
                ready_count += 1
        return ready_count


class BaseAgent(ABC):
    """
    智能体基础抽象类
    
    定义所有白方智能体必须实现的核心接口和通用功能
    """
    
    def __init__(self, agent_name: str = "BaseAgent"):
        """
        初始化智能体
        
        Args:
            agent_name: 智能体名称
        """
        self.agent_name = agent_name
        self.reset_count = 0
        self.step_count = 0
        self.performance_metrics = {}
    
    @abstractmethod
    def get_actions(self, observation: AgentObservation) -> AgentAction:
        """
        根据观测生成动作
        
        Args:
            observation: 当前观测信息
            
        Returns:
            AgentAction: 生成的动作
        """
        pass
    
    @abstractmethod
    def reset(self):
        """重置智能体状态"""
        pass
    
    def update_metrics(self, key: str, value: Any):
        """更新性能指标"""
        self.performance_metrics[key] = value
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取性能指标"""
        return self.performance_metrics.copy()
    
    def step(self, observation: AgentObservation) -> AgentAction:
        """
        执行一步决策
        
        包装了get_actions方法，添加了统计和调试功能
        """
        self.step_count += 1
        action = self.get_actions(observation)
        return action
    
    def get_info(self) -> Dict[str, Any]:
        """获取智能体基本信息"""
        return {
            'name': self.agent_name,
            'reset_count': self.reset_count,
            'step_count': self.step_count,
            'metrics': self.performance_metrics
        }


class RuleBasedAgent(BaseAgent):
    """
    基于规则的智能体基类
    
    为规则AI提供通用的辅助方法和框架
    """
    
    def __init__(self, agent_name: str = "RuleBasedAgent", strategy: str = "balanced"):
        """
        初始化规则智能体
        
        Args:
            agent_name: 智能体名称
            strategy: 策略类型 ("defensive", "aggressive", "balanced")
        """
        super().__init__(agent_name)
        self.strategy = strategy
        self.decision_history = []
        self.target_assignments = {}  # 目标分配记录
    
    def reset(self):
        """重置规则智能体状态"""
        self.reset_count += 1
        self.step_count = 0
        self.decision_history.clear()
        self.target_assignments.clear()
        self.performance_metrics.clear()
    
    def _find_priority_target(self, detected_targets: List['BlackBoat']) -> Optional['BlackBoat']:
        """
        寻找优先目标
        
        Args:
            detected_targets: 探测到的黑方目标列表
            
        Returns:
            优先攻击的目标，如果没有则返回None
        """
        if not detected_targets:
            return None
        
        active_targets = [t for t in detected_targets if t.active and not t.crossed]
        if not active_targets:
            return None
        
        # 按距离突破线的距离排序，距离越近优先级越高
        return min(active_targets, key=lambda t: t.x)
    
    def _calculate_threat_level(self, target: 'BlackBoat') -> float:
        """
        计算目标威胁等级
        
        Args:
            target: 目标黑方无人艇
            
        Returns:
            威胁等级 (0-1)
        """
        from sh15.core.utils import calculate_threat_level
        
        # 计算到突破线的距离
        distance_to_breakthrough = max(0, target.x - 50000)  # A1A6线在x=50000
        
        return calculate_threat_level(distance_to_breakthrough, target.speed)
    
    def _assign_targets(self, white_boats: List['WhiteBoat'], 
                       detected_targets: List['BlackBoat']) -> Dict[int, Optional['BlackBoat']]:
        """
        为白方无人艇分配目标
        
        Args:
            white_boats: 白方无人艇列表
            detected_targets: 探测到的目标列表
            
        Returns:
            目标分配字典 {boat_id: target}
        """
        assignments = {}
        available_targets = detected_targets.copy()
        
        # 按威胁等级排序目标
        available_targets.sort(key=self._calculate_threat_level, reverse=True)
        
        # 按距离为每艘船分配最近的高威胁目标
        for boat in white_boats:
            if not boat.active:
                assignments[boat.id] = None
                continue
            
            if available_targets:
                # 找到距离该船最近的目标
                best_target = min(available_targets, key=lambda t: boat.distance_to(t))
                assignments[boat.id] = best_target
                available_targets.remove(best_target)
            else:
                assignments[boat.id] = None
        
        return assignments
    
    def _record_decision(self, decision_type: str, details: Dict[str, Any]):
        """记录决策信息"""
        decision = {
            'step': self.step_count,
            'type': decision_type,
            'details': details
        }
        self.decision_history.append(decision)
        
        # 限制历史记录长度
        if len(self.decision_history) > 1000:
            self.decision_history = self.decision_history[-500:]


class LearningAgent(BaseAgent):
    """
    学习型智能体基类
    
    为强化学习等学习型AI提供通用接口
    """
    
    def __init__(self, agent_name: str = "LearningAgent"):
        """初始化学习智能体"""
        super().__init__(agent_name)
        self.training_mode = True
        self.model = None
        self.experience_buffer = []
    
    @abstractmethod
    def train(self, experiences: List[Dict[str, Any]]):
        """
        训练智能体
        
        Args:
            experiences: 经验数据列表
        """
        pass
    
    @abstractmethod
    def save_model(self, path: str):
        """保存模型"""
        pass
    
    @abstractmethod
    def load_model(self, path: str):
        """加载模型"""
        pass
    
    def set_training_mode(self, training: bool):
        """设置训练模式"""
        self.training_mode = training
    
    def add_experience(self, experience: Dict[str, Any]):
        """添加经验数据"""
        self.experience_buffer.append(experience)
        
        # 限制缓冲区大小
        if len(self.experience_buffer) > 10000:
            self.experience_buffer = self.experience_buffer[-5000:] 